#  ipv4.rb
#
#  This file is part of myvnet.
#
#  Copyright (C) 2008  Filipe David Borba Manana,  fdmanana@gmail.com
#
#  myvnet is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  myvnet is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with myvnet.  If not, see <http://www.gnu.org/licenses/>.


require 'lib/protocols/generic_protocol'
require 'lib/protocols/icmpv4'
require 'lib/protocols/udp'
require 'lib/exceptions'


module IPv4
    
  # === Synopsis
  #
  #   IPv4.to_friendly_addr(stream)
  #
  # === Args
  #
  # +stream+:: 
  #   A byte stream representing an IPv4 address (only first 4 bytes are processed,
  #   the rest of the stream is ignored).
  #
  # === Return
  #
  # String representing the IPv4 address as a dot separated quartet of decimal values -
  # the human friendly format.
  #
  # === Exceptions
  #
  # +BadArgType+:: 
  #   If stream parameter is not of type String (byte stream).
  # +BadArgFormat+:: 
  #   If length of the stream is less than 4 bytes.
  #
  # === Description
  #
  # Converts a byte stream, containing an IPv4 address, to a string representation in the
  # human friendly format (dotted notation).
  # 
  def self.to_friendly_addr(stream)
    raise BadArgType, "Expected a byte stream" unless stream.is_a? String
    raise BadArgFormat, "Expected a byte stream with at least 4 bytes" unless stream.length >= 4
    
    bytes = stream.unpack 'C4'        
    "%d.%d.%d.%d" % bytes
  end
  
  
  # === Synopsis
  #
  #   IPv4.to_bstream_addr(address)
  #
  # === Args
  #
  # +address+:: 
  #   A string representing an IPv4 address in the human friendly format
  #   (4 decimal values, each with 1 to 3 digits, separated by dots).
  #
  # === Return
  #
  # A byte stream, in network byte order, representation of the given address.
  #
  # === Exceptions
  #
  # +BadArgType+:: 
  #   If stream parameter is not of type String (byte stream).
  # +BadArgFormat+:: 
  #   If the given address is not a quartet of dot separated decimal values
  #   (each with 1 to 3 digits).
  #
  # === Description
  #
  # Converts a string containing an IPv4 address in the human friendly format (4 decimal values
  # separated by dots) to a byte stream in network byte order (big endian).
  #
  def self.to_bstream_addr(address)
    raise BadArgType, "Expected a String" unless address.is_a? String
    
    if !( address =~ /^\s*\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\s*$/ )
      raise BadArgFormat, "Expected an IPv4 dotted address"
    end    
    
    bytes = (address.scan(/\d{1,3}/)).map { |b| b.to_i }        
    bytes.pack 'C4'
  end
  
  
  # maps protocol field codes to symbols
  PROTO_SYM = {
    0x01 => :icmp,
    0x06 => :tcp,
    0x11 => :udp,    
  }.freeze
  
  # maps protocol field symbols to codes
  PROTO_CODE = PROTO_SYM.invert.freeze
  
    
  # IPv4 packet.
  #
  class Packet < GenericProtocol::Packet
    # minimum header length in bytes (bigger if packet has options)
    HLEN = 20
  
    # default Time To Live value
    DEF_TTL = 64
    
    @packet_type = :ipv4
    @layer = 3
    
    attr_accessor :version
    attr_accessor :hlen         # header length (number of 32 bit words)
    attr_accessor :tos          # type of service
    attr_accessor :len          # total length in bytes (IPv4 header + payload)
    attr_accessor :id           # identification
    attr_accessor :flags
    attr_accessor :frag_offset  # fragment offset
    attr_accessor :ttl          # time to live
    attr_accessor :proto        # protocol
    attr_accessor :checksum     # header checksum
    attr_accessor :src          # source address
    attr_accessor :dst          # destination address
    attr_reader   :options
    
    def options=(opts)
      @options = opts
      @hlen = HLEN
      @hlen += opts.length if opts.is_a? String
      @hlen = @hlen >> 2
    end
    
    alias ihl hlen
    alias ihl= hlen=
    
            
    def initialize(bstream = nil)
      super
      
      if bstream.nil?
        bstream = "\x45\x00\x00\x05\x00\x00\x00\x00\x40\x00"
        bstream << "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
      end
      
      byte = (bstream[0,1].unpack 'C')[0]
      @version = (byte & 0xf0) >> 4
      @hlen = byte & 0x0f
      
      @tos = (bstream[1,1].unpack 'C')[0]
      @len = (bstream[2,2].unpack 'n')[0]
      @id = (bstream[4,2].unpack 'n')[0]
      
      short = (bstream[6,2].unpack 'n')[0]
      @flags = (short & 0xe0_00) >> 13
      @frag_offset = short & 0x1f_ff
      
      @ttl = (bstream[8,1].unpack 'C')[0]
      @proto = (bstream[9,1].unpack 'C')[0]
      @proto = PROTO_SYM[@proto] || @proto
      
      @checksum = (bstream[10,2].unpack 'n')[0]
      
      begin
        @src = IPv4.to_friendly_addr(bstream[12,4])
        @dst = IPv4.to_friendly_addr(bstream[16,4])
      rescue
        raise BadArgFormat, "Invalid IPv4 address(es) in the byte stream"
      end
      
      count = (@hlen << 2) - HLEN
      @options = bstream[HLEN, count]
      
      plen = @len - (@hlen << 2)
      @payload = bstream[(@hlen << 2), plen]
    end
    
    
    def initialize_copy(orig)
      super
      @src = orig.src.dup
      @dst = orig.dst.dup
      @options = orig.options.dup
    end
    
    
    def byte_length
      length = super
      length += @options.length
    end
    
    
    def build_upper_layers!
      return self unless payload? and @payload.is_a? String
      
      case @proto
        when :icmp
          @payload = ICMPv4::Packet.new(self, @payload)
        when :udp
          @payload = UDP::Packet.new(@payload)
        else
          return self
      end
      
      @payload.build_upper_layers!
      @payload.parent = self

      self
    end
    

    # === Synopsis
    #
    #   IPv4::Packet#to_bstream
    #
    # === Return
    #
    # A byte stream, in network byte order, representation of the receiver
    # IPv4 packet.
    #
    # === Description
    #
    # Builds a byte stream (String instance), in network byte order, 
    # representation of the receiver IPv4 packet.
    #
    # *Notes*: 
    # * It calculates and sets the length field for the receiving IPv4 packet.
    # * If the receiving packet's checksum is zero, the checksum is calculated
    #   and set.
    #        
    def to_bstream
      payload_stream = ''
      
      case @payload
        when String
          payload_stream << @payload
        when GenericProtocol::Packet
          payload_stream << (@payload.to_bstream)
      end
      
      @len = HLEN + @options.length + payload_stream.length
      
      header_stream + payload_stream
    end
  
  
    # === Synopsis
    # 
    #   IPv4::Packet#checksum_calc!
    #   
    # === Return
    # 
    # The checksum that was computed for the IPv4 packet's header.
    # 
    # === Description
    # 
    # Computes the checksum of an IPv4 packet's header and sets the
    # receiver packet's checksum field value to the computed value.
    # 
    # *Note*: it also computes the packet's length field value and
    # sets it to that value.
    #
    def checksum_calc!
      self.len_calc!
      @checksum = 0
      header_stream
      @checksum
    end
    

    # === Synopsis
    # 
    #   IPv4::Packet#len_calc!
    #   
    # === Return
    # 
    # The computed length value for the packet's length field.
    # 
    # === Description
    # 
    # Computes the length of an IPv4 packet and sets the receiver
    # packet's length field to the computed value.
    # 
    def len_calc!
      @len = self.byte_length
    end
    

    # === Synopsis
    #
    #   IPv4::Packet#has_tcp?
    #
    # === Return
    #
    # True if this packet encapsulates a TCP packet, false otherwise.
    #
    # === Description
    #
    # Verifies if this packet encapsulates a TCP packet.
    # 
    # *Note*: it invokes IPv4::Packet#build_upper_layers! on self.
    #
    def has_tcp?
      self.build_upper_layers! if @payload.is_a? String
      
      @proto == :tcp
    end
    

    # === Synopsis
    #
    #   IPv4::Packet#has_udp?
    #
    # === Return
    #
    # True if this packet encapsulates a UDP packet, false otherwise.
    #
    # === Description
    #
    # Verifies if this packet encapsulates a UDP packet.
    # 
    # *Note*: it invokes IPv4::Packet#build_upper_layers! on self.
    #    
    def has_udp?
      self.build_upper_layers! if @payload.is_a? String
      
      @proto == :udp
    end
    
    
    # === Synopsis
    #
    #   IPv4::Packet#has_icmp?
    #
    # === Return
    #
    # True if this packet encapsulates an ICMP (v4) packet, false otherwise.
    #
    # === Description
    #
    # Verifies if this packet encapsulates an ICMP (v4) packet.
    # 
    # *Note*: it invokes IPv4::Packet#build_upper_layers! on self.
    #    
    def has_icmp?
      self.build_upper_layers! if @payload.is_a? String
      
      @proto == :icmp
    end
    

    # === Synopsis
    #
    #   IPv4::Packet#is_fragment?
    #
    # === Return
    #
    # True if this IPv4 packet is a fragment.
    #
    # === Description
    #
    # Verifies if this IPv4 packet is a fragment.
    #        
    def is_fragment?
      return true if @frag_offset > 0 or ( @flags & 0b001 ) == 0b001
      false
    end
    
  
    def to_s(indent_level = 0)
      prefix = ' ' * indent_level
      str = ''
      str << prefix << "#{self.class.name}\n"
      str << prefix << "version: #{@version}, hlen: #{@hlen}, tos: #{@tos}, "
      str << "len: #{@len}, id: #{@id}, "
      str << "flags: 0b" << ("%03b" % @flags) << ",\n"
      str << prefix << "offset: #{@frag_offset}, ttl: #{@ttl}, proto: "
      
      case @proto
        when Numeric
          str << "0x" << ("%04X" % @proto)
        else
          str << "#{@proto.to_s}"
      end
      
      str << ", chksum: 0x" << ("%04X" % @checksum) << ",\n"
      str << prefix << "src: #{@src}, dst: #{@dst},\n"
      if not @options.length.zero?
        str << prefix << "options: " << hexdump(@options) << ",\n"
      end
      str << prefix << "payload:\n"
      
      case @payload
        when String
          str << prefix << '   ' << hexdump(@payload) << "\n"
        when GenericProtocol::Packet
          str << @payload.to_s(indent_level + 3)
      end
      
      str
    end    
    
    
    private
    
    def header_stream
      hdr = ''
      
      byte = (@version << 4) | @hlen
      hdr << ( [byte].pack 'C' )      
      hdr << ( [@tos].pack 'C' )
      hdr << ( [@len].pack 'n' )
      hdr << ( [@id].pack 'n' )
      
      short = (@flags << 13) | @frag_offset
      hdr << ( [short].pack 'n' )
      hdr << ( [@ttl].pack 'C' )
      
      if @proto.is_a? Symbol
        hdr << ( [ PROTO_CODE[@proto] ].pack 'C' )
      else
        hdr << ( [ @proto ].pack 'C' )
      end
      
      hdr << ( [@checksum].pack 'n' )
      hdr << IPv4.to_bstream_addr(@src)
      hdr << IPv4.to_bstream_addr(@dst)      
      hdr << @options
      
      set_checksum hdr if @checksum.zero?
      
      hdr
    end
    
        
    def set_checksum(hdr_stream)
      # if odd number of bytes, add an extra one
      padded = false
      if hdr_stream.length[0] == 1
        hdr_stream << "\x00"
        padded = true
      end
      
      hdr_stream[10,2] = "\x00\x00"
      words = hdr_stream.unpack 'n*'

      sum = words.inject do |w,a|
        w + a
      end
      
      sum = (sum & 0xff_ff) + (sum >> 16)
      chksum = (~sum) & 0xff_ff
      @checksum = chksum
      
      hdr_stream[10,2] = [chksum].pack 'n'
      hdr_stream[-1,1] = '' if padded
      
      hdr_stream
    end
        
  end
  
end
